from image_augmentationTSL import resizeImg,random_flip,cropAngle,rotate_image,random_hsv_transform,random_gamma_transform
import numpy as np
import cv2,os,copy
from videoLabel import videoLabel as vl

def showImg(src,locs):
    # img = cv2.imread(src)
    for loc in locs:
        x,y,w,h = map(int,loc)
        cv2.rectangle(src,(x,y),(x+w,y+h),(0,0,255),2)
    cv2.imshow("img",src)
    cv2.waitKey()

def cvtResize(loc,ratResize):
    return [i*ratResize for i in loc]

def autolabelsizeRes(src="",dst="",randomSize=1):
    result = cv2.imread(src)
    _,iw,ih = result.shape[::-1]
    resize_img = resizeImg(int(iw * randomSize), int(ih * randomSize), result)
    newLocs = []
    with open(src.replace(".jpg",".txt"),"r") as f:
        lines = f.readlines()
        for line in lines:
            line = line.rstrip()
            label,x,y,w,h = map(float,line.split(" "))
            print(label,x,y,w,h)
            rx,ry,rw,rh = vl.convertR(x,y,w,h,iw,ih)
            print(rx,ry,rw,rh)
            newX,newY,newW,newH = cvtResize([rx,ry,rw,rh],randomSize)
            newLocs.append([label,newX, newY, newW, newH])

    return newLocs,resize_img

def autoFlip(img,locs,model):
    new_loc = []
    if model==0:
        result = random_flip(img, 1)  # model: 水平镜像1,垂直镜像0,对角镜像-1
        _,iw,ih = result.shape[::-1]
        for loc in locs:
            label,x,y,w,h = loc
            newX = iw - x - w
            new_loc.append([label,newX,y,w,h])
        return new_loc,result
    else:
        return locs,img

def findLoc(img):
    # locs = []
    gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
    _,thresh = cv2.threshold(gray,200,255,cv2.THRESH_BINARY)
    cnts,_ = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_NONE)
    locs = [list(cv2.boundingRect(cnt)) for cnt in cnts]
    if(len(locs)==0):
        return [None]
    return locs[0]

def rotateImg(img, angle, crop,locs):
    new_loc = []

    img_rotated = rotate_image(img, angle, crop)
    for loc in locs:
        back = np.zeros_like(img)
        label,x,y,w,h = map(int ,loc)
        back[slice(y,y+h),slice(x,x+w)]  = (255,255,255)
        cv2.imshow("back",back)
        cv2.waitKey(1)
        back_rotated=rotate_image(back, angle, crop)
        back_locs = findLoc(back_rotated)
        try:
            rx,ry,rw,rh = back_locs
        except:
            rx,ry,rw,rh = None,None,None,None
        new_loc.append([label,rx,ry,rw,rh])

    return  new_loc,img_rotated


def auto(src="",dst=""):
    img = cv2.imread(src)
    _,iw,ih = img.shape[::-1]

    ####################### auto resize #######################
    randomSize = np.random.randint(8, 13) / 10.0
    resize_loc,resize_img = autolabelsizeRes(src,"",randomSize)
    # showImg(resize_img,resize_loc)

    ####################### auto flip #######################
    ifFlip = np.random.randint(0, 2)
    flip_loc,flip_img = autoFlip(resize_img, resize_loc, ifFlip)
    # showImg(flip_img,new_loc)

    ####################### auto rota #######################
    angle,crop = cropAngle(15)
    crop = False
    new_loc,img_rotated = rotateImg(flip_img, angle, crop, flip_loc)
    # showImg(img_rotated,new_loc)

    ####################### auto hsv #######################
    hsv_img = random_hsv_transform(img_rotated, 30, 0.1, 0.20)

    ####################### auto gama #######################
    game_img = random_gamma_transform(hsv_img, 0.8)

    return new_loc,game_img


def main():
    src = "/mnt/data/people/falldown"
    dst = "/mnt/data/people/falldown_auto"
    labels = "/mnt/data/people/falldown_auto"
    toWrite = True
    for i in os.listdir(src):
        if(i.endswith(".png")):
            name = i.split(".")[0]
            for num in range(10):
                new_loc,img = auto(os.path.join(src,i))
                save_img = copy.copy(img)
                _,iw,ih = save_img.shape[::-1]

                for loc in new_loc:
                    label,x,y,w,h = loc
                    if x != None:
                        cv2.rectangle(img,(x,y),(x+w,y+h),(0,0,255),2)
                        toWrite=True
                    else:
                        toWrite=False
                    # cv2.imwrite("data/tmp/"+i,img)

                cv2.imwrite(os.path.join(dst, name + "_" + str(num) + ".png"), save_img)
                with open(os.path.join(labels,name + "_" + str(num) + ".txt"),'w') as f:
                    for loc in new_loc:
                        label,x, y, w, h = loc
                        if(x!=None):
                            x1,x2,y1,y2 = x,x+w,y,y+h
                            rx,ry,rw,rh = vl.convert([iw,ih],[x1,x2,y1,y2])
                            r_loc = ["{:.6f}".format(i) for i in [rx, ry, rw, rh]]
                            r_loc.insert(0, str(label))
                            f.writelines( " ".join(map(str,r_loc))+"\n" )


if __name__ == "__main__":
    main()